#!/usr/bin/env python
# encoding: utf-8
# @Time : 2020/6/22 19:02 
# @Author : 能量咖啡豆 
# @File : lwlr.py 
# @desc : lwlr.py 

import numpy as np
import matplotlib.pyplot as plt

"""
# 读取文件
# dataMat输入变量
# labelMat输出变量
"""
def loadDataSet(fileName):
    numFeat = len(open(fileName).readline().split('\t')) - 1
    dataMat = []
    labelMat = []

    fr = open(fileName)
    for line in fr.readlines():
        lineArr = []
        curLine = line.strip().split('\t')
        for i in range(numFeat):
            lineArr.append(float(curLine[i]))
        dataMat.append(lineArr)
        labelMat.append(float(curLine[-1]))
    return dataMat, labelMat

def lwlr(testPoint, xArr, yArr, k=1.0):
    xMat = np.mat(xArr)
    yMat = np.mat(yArr).T
    m = np.shape(xMat)[0]               #m为xMat的行数
    weights = np.mat(np.eye((m)))
    for j in range(m):
        diffMat = testPoint - xMat[j,:]
        weights[j,j] = np.exp(diffMat*diffMat.T/(-2.0*k**2))
    xTx = xMat.T * (weights * xMat)
    if np.linalg.det(xTx) == 0.0:
        print("This Matrix is singular, cannot do inverse")
        return
    ws = xTx.I * (xMat.T * (weights * yMat))
    return testPoint * ws

def lwlrTest(testArr, xArr, yArr, k=1.0):
    m = np.shape(testArr)[0]
    yHat = np.zeros(m)
    for i in range(m):
        yHat[i] = lwlr(testArr[i], xArr, yArr, k)
    return yHat

if __name__ == "__main__":
    print("lwlr局部加权线性回归")
    xArr, yArr = loadDataSet('data/abalone.txt')
    yHat = lwlrTest(xArr, xArr, yArr, 0.01)
    xMat = np.mat(xArr)

    #这块计算有点复杂，需要进一步学习
    srtInd = xMat[:,1].argsort(0)
    xSort = xMat[srtInd][:,0,:]

    #绘图
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.plot(xSort[:,1],yHat[srtInd])
    ax.scatter(xMat[:,1].flatten().A[0],np.mat(yArr).T.flatten().A[0], s=2, c='red')
    plt.show()